#!/usr/bin/env python3
# @Date    : 2020-07-14
# @Author  : Bright Li (brt2@qq.com)
# @Link    : https://gitee.com/brt2

# import imageio
from PIL import Image, ImageDraw, ImageFont
from BlindWatermark.BlindWatermark import watermark
import jhead

_path_font = "static/FreeMono.ttf"
_path_wm = "tmp/wm.jpg"

random_seed1,random_seed2,mod1 = 2743, 6989, 32
block_shape = 2  # 8
dwt_deep = 1


def make_wm_img(text, mask_height):
    mask_width = int(mask_height * 0.6 * len(text)) +1
    mask = Image.new('L', (mask_width, mask_height), 0)

    font = ImageFont.truetype(_path_font, mask_height)  # 更改字体尺寸
    ImageDraw.Draw(mask).text((0,0), fill=255, text=text, font=font)
    mask.save(_path_wm)
    return mask_width, mask_height

def make_wm_img2(text, mask_width, mask_height):
    mask = Image.new('L', (mask_width, mask_height), 0)
    font_size = int(mask_width / len(text) / 0.6)
    font = ImageFont.truetype(_path_font, font_size)  # 更改字体尺寸
    ImageDraw.Draw(mask).text((0,0), fill=255, text=text, font=font)
    mask.save(_path_wm)

def paste_text(path_img, text, path_save, level=3, rotate=30):
    img = Image.open(path_img)

    w, h = img.size
    mask_height = w // (level * 4)

    mask_width = int(mask_height * 0.6 * len(text)) +1
    mask = Image.new('L', (mask_width, mask_height), 0)

    font = ImageFont.truetype(_path_font, mask_height)  # 更改字体尺寸
    ImageDraw.Draw(mask).text((0,0), fill=255, text=text, font=font)

    pos_left = 5
    pos_right = w - mask_width - pos_left
    pos_top = 5
    gap_height = (h - pos_top * 2) // (level * 2)

    for i in range(level):
        pos_height = pos_top + i * gap_height * 2
        img.paste(mask, (pos_left, pos_height), mask)
        pos_height += gap_height
        img.paste(mask, (pos_right, pos_height), mask)

    img.save(path_save)


def paste_bwm(path_img, text, path_save):
    """ 阿里水印 """
    bwm = watermark(random_seed1,
                    random_seed2,
                    mod1,
                    block_shape=(block_shape, block_shape),
                    dwt_deep=dwt_deep)

    bwm.read_ori_img(path_img)
    font_size = 32
    mask_width, mask_height = make_wm_img(text, font_size)
    bwm.read_wm(_path_wm)
    bwm.embed(path_save)
    jhead.set_comment(path_save, "{}x{}".format(mask_height, mask_width))

def check_bwm(path_img, path_extract):
    exif_comment = jhead.get_comment(path_img)
    if not exif_comment:
        print(">> 解密失败，未获取到exif信息")
        return

    wm_shape = [int(i) for i in exif_comment.split("x")]

    bwm = watermark(random_seed1,
                    random_seed2,
                    mod1,
                    wm_shape=wm_shape,
                    block_shape=(block_shape, block_shape),
                    dwt_deep=dwt_deep)
    bwm.extract(path_img, path_extract)


if __name__ == "__main__":
    def getopt():
        import argparse

        parser = argparse.ArgumentParser("Extract-Blind-watermark", description="提取图像的盲水印信息")
        parser.add_argument("path_img", action="store", help="待处理图像路径")
        parser.add_argument("path_output", action="store", help="生成包含信息的水印图像")
        return parser.parse_args()

    args = getopt()
    check_bwm(args.path_img, args.path_output)
    print(">>> 已完成，请查看水印图像：", args.path_output)
